import * as jose from 'jose';
import { InferType } from 'yup';
import { accessTokenPayloadSchema } from './schema-fields';
import { StackAssertionError, throwErr } from "./utils/errors";
import { Store } from "./utils/stores";


export type AccessTokenPayload = InferType<typeof accessTokenPayloadSchema>;

function decodeAccessTokenIfValid(token: string): AccessTokenPayload | null {
  try {
    const payload = jose.decodeJwt(token);
    return accessTokenPayloadSchema.validateSync(payload);
  } catch (e) {
    return null;
  }
}

export class AccessToken {
  static createIfValid(token: string): AccessToken | null {
    const payload = decodeAccessTokenIfValid(token);
    if (!payload) return null;
    return new AccessToken(token);
  }

  private constructor(
    public readonly token: string,
  ) {
    if (token === "undefined") {
      throw new StackAssertionError("Access token is the string 'undefined'; it's unlikely this is the correct value. They're supposed to be unguessable!");
    }
  }

  get payload() {
    return decodeAccessTokenIfValid(this.token) ?? throwErr("Invalid access token in payload (should've been validated in createIfValid)", { token: this.token });
  }

  get expiresAt(): Date {
    const { exp } = this.payload;
    if (exp === undefined) return new Date(8640000000000000);  // max date value
    return new Date(exp * 1000);
  }

  /**
   * @returns The number of milliseconds until the access token expires, or 0 if it has already expired.
   */
  get expiresInMillis(): number {
    return Math.max(0, this.expiresAt.getTime() - Date.now());
  }

  isExpired(): boolean {
    return this.expiresInMillis <= 0;
  }
}

export class RefreshToken {
  constructor(
    public readonly token: string,
  ) {
    if (token === "undefined") {
      throw new StackAssertionError("Refresh token is the string 'undefined'; it's unlikely this is the correct value. They're supposed to be unguessable!");
    }
  }
}

/**
 * An InternalSession represents a user's session, which may or may not be valid. It may contain an access token, a refresh token, or both.
 *
 * A session never changes which user or session it belongs to, but the tokens in it may change over time.
 */
export class InternalSession {
  /**
  * Each session has a session key that depends on the tokens inside. If the session has a refresh token, the session key depends only on the refresh token. If the session does not have a refresh token, the session key depends only on the access token.
  *
  * Multiple Session objects may have the same session key, which implies that they represent the same session by the same user. Furthermore, a session's key never changes over the lifetime of a session object.
  *
  * This is useful for caching and indexing sessions.
  */
  public readonly sessionKey: string;

  /**
   * An access token that is not known to be invalid (ie. may be valid, but may have expired).
   */
  private _accessToken: Store<AccessToken | null>;
  private readonly _refreshToken: RefreshToken | null;

  /**
   * Whether the session as a whole is known to be invalid (ie. both access and refresh tokens are invalid). Used as a cache to avoid making multiple requests to the server (sessions never go back to being valid after being invalidated).
   *
   * It is possible for the access token to be invalid but the refresh token to be valid, in which case the session is
   * still valid (just needs a refresh). It is also possible for the access token to be valid but the refresh token to
   * be invalid, in which case the session is also valid (eg. if the refresh token is null because the user only passed
   * in an access token, eg. in a server-side request handler).
   */
  private _knownToBeInvalid = new Store<boolean>(false);

  private _refreshPromise: Promise<AccessToken | null> | null = null;

  constructor(private readonly _options: {
    refreshAccessTokenCallback(refreshToken: RefreshToken): Promise<AccessToken | null>,
    refreshToken: string | null,
    accessToken?: string | null,
  }) {
    this._accessToken = new Store(_options.accessToken ? AccessToken.createIfValid(_options.accessToken) : null);
    this._refreshToken = _options.refreshToken ? new RefreshToken(_options.refreshToken) : null;
    if (_options.accessToken === null && _options.refreshToken === null) {
      // this session is already invalid
      this._knownToBeInvalid.set(true);
    }
    this.sessionKey = InternalSession.calculateSessionKey({ accessToken: _options.accessToken ?? null, refreshToken: _options.refreshToken });
  }

  static calculateSessionKey(ofTokens: { refreshToken: string | null, accessToken?: string | null }): string {
    if (ofTokens.refreshToken) {
      return `refresh-${ofTokens.refreshToken}`;
    } else if (ofTokens.accessToken) {
      return `access-${ofTokens.accessToken}`;
    } else {
      return "not-logged-in";
    }
  }

  isKnownToBeInvalid() {
    return this._knownToBeInvalid.get();
  }

  /**
   * Marks the session object as invalid, meaning that the refresh and access tokens can no longer be used.
   */
  markInvalid() {
    this._accessToken.set(null);
    this._knownToBeInvalid.set(true);
  }

  onInvalidate(callback: () => void): { unsubscribe: () => void } {
    return this._knownToBeInvalid.onChange(() => callback());
  }

  /**
   * Returns the access token if it is found in the cache and not expired yet, or null otherwise. Never fetches new tokens.
   */
  getAccessTokenIfNotExpiredYet(minMillisUntilExpiration: number): AccessToken | null {
    if (minMillisUntilExpiration > 60_000) {
      throw new Error(`Required access token expiry ${minMillisUntilExpiration}ms is too long; access tokens are too short to be used for more than 60s`);
    }

    const accessToken = this._getPotentiallyInvalidAccessTokenIfAvailable();
    if (!accessToken || accessToken.expiresInMillis < minMillisUntilExpiration) return null;
    return accessToken;
  }

  /**
   * Returns the access token if it is found in the cache, fetching it otherwise.
   *
   * This is usually the function you want to call to get an access token. Either set `minMillisUntilExpiration` to a reasonable value, or catch errors that occur if it expires, and call `markAccessTokenExpired` to mark the token as expired if so (after which a call to this function will always refetch the token).
   *
   * @returns null if the session is known to be invalid, cached tokens if they exist in the cache and the access token hasn't expired yet (the refresh token might still be invalid), or new tokens otherwise.
   */
  async getOrFetchLikelyValidTokens(minMillisUntilExpiration: number): Promise<{ accessToken: AccessToken, refreshToken: RefreshToken | null } | null> {
    const accessToken = this.getAccessTokenIfNotExpiredYet(minMillisUntilExpiration);
    if (!accessToken) {
      const newTokens = await this.fetchNewTokens();
      const expiresInMillis = newTokens?.accessToken.expiresInMillis;
      if (expiresInMillis && expiresInMillis < minMillisUntilExpiration) {
        throw new StackAssertionError(`Required access token expiry ${minMillisUntilExpiration}ms is too long; access tokens are too short when they're generated (${expiresInMillis}ms)`);
      }
      return newTokens;
    }
    return { accessToken, refreshToken: this._refreshToken };
  }

  /**
   * Fetches new tokens that are, at the time of fetching, guaranteed to be valid.
   *
   * The newly generated tokens are short-lived, so it's good practice not to rely on their validity (if possible). However, this function is useful in some cases where you only want to pass access tokens to a service, and you want to make sure said access token has the longest possible lifetime.
   *
   * In most cases, you should prefer `getOrFetchLikelyValidTokens`.
   *
   * @returns null if the session is known to be invalid, or new tokens otherwise (which, at the time of fetching, are guaranteed to be valid).
   */
  async fetchNewTokens(): Promise<{ accessToken: AccessToken, refreshToken: RefreshToken | null } | null> {
    const accessToken = await this._getNewlyFetchedAccessToken();
    return accessToken ? { accessToken, refreshToken: this._refreshToken } : null;
  }

  markAccessTokenExpired(accessToken: AccessToken) {
    // TODO we don't need this anymore, since we now check the expiry by ourselves
    if (this._accessToken.get() === accessToken) {
      this._accessToken.set(null);
    }
  }

  /**
   * Note that a callback invocation with `null` does not mean the session has been invalidated; the access token may just have expired. Use `onInvalidate` to detect invalidation.
   */
  onAccessTokenChange(callback: (newAccessToken: AccessToken | null) => void): { unsubscribe: () => void } {
    return this._accessToken.onChange(callback);
  }

  /**
   * @returns An access token, which may be expired or expire soon, or null if it is known to be invalid.
   */
  private _getPotentiallyInvalidAccessTokenIfAvailable(): AccessToken | null {
    if (!this._refreshToken) return null;
    if (this.isKnownToBeInvalid()) return null;

    const accessToken = this._accessToken.get();
    if (accessToken && !accessToken.isExpired()) return accessToken;

    return null;
  }

  /**
   * You should prefer `_getOrFetchPotentiallyInvalidAccessToken` in almost all cases.
   *
   * @returns A newly fetched access token (never read from cache), or null if the session either does not represent a user or the session is invalid.
   */
  private async _getNewlyFetchedAccessToken(): Promise<AccessToken | null> {
    if (!this._refreshToken) return null;
    if (this._knownToBeInvalid.get()) return null;

    if (!this._refreshPromise) {
      this._refreshAndSetRefreshPromise(this._refreshToken);
    }
    return await this._refreshPromise;
  }

  private _refreshAndSetRefreshPromise(refreshToken: RefreshToken) {
    let refreshPromise: Promise<AccessToken | null> = this._options.refreshAccessTokenCallback(refreshToken).then((accessToken) => {
      if (refreshPromise === this._refreshPromise) {
        this._refreshPromise = null;
        this._accessToken.set(accessToken);
        if (!accessToken) {
          this.markInvalid();
        }
      }
      return accessToken;
    });
    this._refreshPromise = refreshPromise;
  }
}
